'''
Created on Feb 19, 2013

@author: nzhao
'''

import numpy as np
from traits.api import HasTraits, Int, List

class QMobject(HasTraits):
    nbody      = Int()
    dim_list   = List()
    basis_list = List()
    dim        = Int()
    
    def __init__(self, nbody, dim_list):
        self.nbody    = nbody
        self.dim_list = dim_list
        self.dim      = int( np.prod(self.dim_list) )
        
        self.basis_list = self._generate_basis(self.dim_list, self.dim)

    def _generate_basis(self, dim_list, dim):
        res = [ list(self._index2basis(i, dim_list))[::-1] for i in range(dim) ]
        return res
    
    def _index2basis(self, x, lst):
        r = 0
        for i in lst:
            r =  x % i
            x = (x - r) / i
            yield r
    
    def _basis2index(self, basis, dim_lst):
        acc_dim_lst = np.array([np.prod(dim_lst[:k]) for k in range(len(dim_lst)) ])
        return int( sum(acc_dim_lst * np.array(basis) ) )
